{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8e26a06e-e42e-43fd-9d87-fecc1569a9ec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-03-07 06:39:25,370] [INFO] [real_accelerator.py:161:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "初始化设置\n",
    "\"\"\"\n",
    "import os\n",
    "import io\n",
    "import re\n",
    "import ast\n",
    "import pdb\n",
    "import clip\n",
    "import json\n",
    "import torch\n",
    "import faiss\n",
    "import chromadb\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from tqdm import tqdm\n",
    "Image.MAX_IMAGE_PIXELS = None\n",
    "import matplotlib.pyplot as plt\n",
    "from transformers import AutoModel, AutoTokenizer\n",
    "from transformers import CLIPProcessor, CLIPModel\n",
    "\n",
    "### 初始化一个clip模型\n",
    "clip_model, preprocess = clip.load(\"ViT-B/16\", device=\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "486171f9-b085-4a40-8725-6da6114749bd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "196\n"
     ]
    }
   ],
   "source": [
    "# line_mapping = {}\n",
    "# file_path = \"/mnt/petrelfs/liuziyu/LLM_Memory/SimplyRetrieve/classification/New_test/database/cars196/classnames.txt\"\n",
    "# with open(file_path, 'r') as file:\n",
    "#     for line_number, line in enumerate(file):\n",
    "#         line_mapping[line_number] = line.strip()  \n",
    "# print(len(line_mapping))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36b567ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "Birds\n",
    "top_k =2 \n",
    "top_number = 3\n",
    "Clustering ACC: 51.632723507076285\n",
    "Semantic ACC:   69.47120666503906\n",
    "************************************\n",
    "Pets\n",
    "top_k = 3\n",
    "top_number = 2\n",
    "Clustering ACC: 74.12101390024529\n",
    "Semantic ACC:   74.83729553222656\n",
    "************************************\n",
    "Flowers\n",
    "top_k = 3\n",
    "top_number = 1\n",
    "Clustering ACC: 63.69653602211742\n",
    "Semantic ACC:   53.238861083984375\n",
    "************************************\n",
    "Cars\n",
    "top_k = 2\n",
    "top_number = 1\n",
    "Clustering ACC: 53.156323840318365\n",
    "Semantic ACC:   63.64885711669922\n",
    "\n",
    "************************************\n",
    "Dogs\n",
    "top_k = 2\n",
    "top_number = 3\n",
    "Clustering ACC: 50.00815850815851\n",
    "Semantic ACC:   65.2044677734375\n",
    "************************************"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aba4050f-3b24-4203-a5ae-3f0902214ed5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "\n",
    "### best result record above, change top_k and top_number first.\n",
    "top_k = 2\n",
    "top_number = 3\n",
    "\n",
    "index_img_save_path =  \"/mnt/petrelfs/liuziyu/LLM_Memory/SimplyRetrieve/classification/New_test/database/cub200/cub200_database.index\"\n",
    "index = faiss.read_index(index_img_save_path)\n",
    "\n",
    "image_paths = []\n",
    "gt_labels = []\n",
    "val_file_path = \"/mnt/petrelfs/liuziyu/LLM_Memory/SimplyRetrieve/classification/New_test/database/cub200/cub200_database.txt\"\n",
    "with open(val_file_path, 'r') as file:\n",
    "    for line in file:\n",
    "        image_path = line.split(' ')[0]\n",
    "        label = line.split(' ')[1].strip()\n",
    "        image_paths.append(image_path)\n",
    "        gt_labels.append(label)\n",
    "\n",
    "selected_names = []\n",
    "\n",
    "trainset_file_path = \"/mnt/petrelfs/liuziyu/LLM_Memory/SimplyRetrieve/classification/New_test/database/cub200/paired_data_birds.txt\"\n",
    "with open(trainset_file_path, 'r') as file:\n",
    "    lines = file.readlines()\n",
    "    for idx, image_path in enumerate(tqdm(image_paths, desc=\"Process:\")):\n",
    "        # print(image_path)\n",
    "        # print(gt_labels[idx])\n",
    "        # print(line_mapping[int(gt_labels[idx])])\n",
    "    \n",
    "        with torch.no_grad():\n",
    "            image = preprocess(Image.open(image_path)).unsqueeze(0).to(\"cuda\")\n",
    "            # torch.Size([1, 512])\n",
    "            image_features = clip_model.encode_image(image)\n",
    "            image_features /= image_features.norm(dim=-1, keepdim=True)\n",
    "            image_features = np.array(image_features.cpu())\n",
    "            distance, index_result = index.search(image_features, top_k)\n",
    "\n",
    "            retrieved_names = []\n",
    "            for index_number in index_result[0]:\n",
    "                parts = lines[index_number].strip()\n",
    "                list_part = parts.split('[')[-1].rstrip(']')\n",
    "                list_extracted = ast.literal_eval('[' + list_part + ']')\n",
    "                retrieved_names.extend(list_extracted[:top_number])\n",
    "            # print(retrieved_names)\n",
    "\n",
    "            # most_common_element = Counter(retrieved_names).most_common(1)[0][0]\n",
    "            # print(most_common_element)\n",
    "            # selected_names.append(most_common_element)\n",
    "\n",
    "            selected_names.extend(retrieved_names)\n",
    "            \n",
    "print(len(selected_names))\n",
    "print(selected_names)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
